import torch
import cv2
import numpy as np

class YOLOSubjectDetector:
    """
    使用YOLO模型检测图像中的主体，并返回其边界框的实现类。
    """

    def __init__(self,experiments_path,logger, conf_threshold=0.5, iou_threshold=0.4, device="npu"):
        """
        初始化检测器。
        
        参数:
        - model_path: 预训练YOLO模型的路径。
        - conf_threshold: 置信度阈值，用于过滤低置信度的检测结果。
        - iou_threshold: 非极大值抑制的IOU阈值，用于移除多余的框。
        - device: 运行设备npu。
        """
        self.device = device
        self.conf_threshold = conf_threshold
        self.iou_threshold = iou_threshold
        logger.info("Loading YOLO model...")
        # 加载YOLO模型
        self.model = torch.hub.load('ultralytics/yolov5', 'custom', path=experiments_path).to(self.device)
        self.model.eval()  # 设置为评估模式
        logger.info(f"YOLO model loaded successfully. Device: {self.device}")
        # 记录参数
        logger.info(f"Confidence Threshold: {self.conf_threshold}, IOU Threshold: {self.iou_threshold}")
        self.logger = logger
        
    def preprocess_image(self, image_path):
        """
        读取并预处理输入图像。

        参数:
        - image_path: 图像文件的路径。

        返回:
        - image: 原始图像，用于可视化。
        - input_tensor: 预处理后的图像张量，用于模型输入。
        """
        # 使用OpenCV读取图像（BGR格式）
        image = cv2.imread(image_path)
        # 转换为RGB格式
        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        # 使用YOLO的默认预处理（缩放到640x640，归一化等）
        input_tensor = self.model.model.preprocess(image_rgb)
        return image, input_tensor

    def detect(self, image_path):
        """
        检测图像中的主体并返回主体的边界框。

        参数:
        - image_path: 输入图像的路径。

        返回:
        - result_bbox: 包含主体的边界框 [x_min, y_min, x_max, y_max]。
        """
        # 读取并预处理图像
        image, input_tensor = self.preprocess_image(image_path)
        # 模型推理
        results = self.model(input_tensor)
        # 获取检测结果 (包括 bbox, 置信度, 类别)
        detections = results.pred[0].cpu().numpy()  # 转为NumPy格式
        
        # 筛选出高置信度的检测结果
        filtered_boxes = [
            detection[:4]  # bbox 的坐标 [x_min, y_min, x_max, y_max]
            for detection in detections
            if detection[4] >= self.conf_threshold  # 置信度过滤
        ]
        
        # 如果没有满足条件的框，返回None
        if not filtered_boxes:
            self.logger.warning("未检测到主体。")
            return None
        
        # 根据面积选择最大的框作为主体框
        main_bbox = max(filtered_boxes, key=lambda box: (box[2] - box[0]) * (box[3] - box[1]))
        return main_bbox

    def visualize(self, image_path, bbox,save_path=None):
        """
        在图像上绘制边界框并显示。

        参数:
        - image_path: 输入图像的路径。
        - bbox: 检测到的边界框 [x_min, y_min, x_max, y_max]。
        """
        # 读取原始图像
        image = cv2.imread(image_path)
        # 将边界框绘制到图像上
        x_min, y_min, x_max, y_max = map(int, bbox)
        cv2.rectangle(image, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)  # 绿色框
        cv2.putText(image, "Subject", (x_min, y_min - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
        # 显示图像
        if save_path:
            cv2.imwrite(save_path, image)
            self.logger.info(f"检测结果已保存至: {save_path}")

# 示例代码
if __name__ == "__main__":
    # 创建检测器实例
    detector = YOLOSubjectDetector(model_path="yolov5s.pt", conf_threshold=0.5, device="cpu")

    # 输入图像路径
    image_path = "dog_on_grass.jpg"

    # 检测主体
    bbox = detector.detect(image_path)

    if bbox:
        print(f"检测到的主体边界框: {bbox}")
        # 可视化检测结果
        detector.visualize(image_path, bbox)
